"""
统计函数模块，实现各种统计计算功能。
"""
import torch
import numpy as np
from typing import Dict, List, Optional, Union, Tuple, Any

def calculate_norm(tensor: torch.Tensor) -> float:
    """
    计算张量的L2范数
    
    Args:
        tensor: 输入张量
    
    Returns:
        L2范数值
    """
    return torch.norm(tensor).item()

def calculate_max(tensor: torch.Tensor) -> float:
    """
    计算张量的最大值
    
    Args:
        tensor: 输入张量
    
    Returns:
        最大值
    """
    return torch.max(tensor).item()

def calculate_min(tensor: torch.Tensor) -> float:
    """
    计算张量的最小值
    
    Args:
        tensor: 输入张量
    
    Returns:
        最小值
    """
    return torch.min(tensor).item()

def calculate_mean(tensor: torch.Tensor) -> float:
    """
    计算张量的均值
    
    Args:
        tensor: 输入张量
    
    Returns:
        均值
    """
    return torch.mean(tensor).item()

def calculate_std(tensor: torch.Tensor) -> float:
    """
    计算张量的标准差
    
    Args:
        tensor: 输入张量
    
    Returns:
        标准差
    """
    return torch.std(tensor).item()

def calculate_zeros(tensor: torch.Tensor) -> int:
    """
    计算张量中零元素的数量
    
    Args:
        tensor: 输入张量
    
    Returns:
        零元素数量
    """
    return (tensor == 0).sum().item()

def calculate_nans(tensor: torch.Tensor) -> int:
    """
    计算张量中NaN元素的数量
    
    Args:
        tensor: 输入张量
    
    Returns:
        NaN元素数量
    """
    return torch.isnan(tensor).sum().item()

def calculate_infs(tensor: torch.Tensor) -> int:
    """
    计算张量中Inf元素的数量
    
    Args:
        tensor: 输入张量
    
    Returns:
        Inf元素数量
    """
    return torch.isinf(tensor).sum().item()

def calculate_histogram(tensor: torch.Tensor, bins: int = 10) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    计算张量的直方图
    
    Args:
        tensor: 输入张量
        bins: 直方图箱数
    
    Returns:
        直方图统计和边界
    """
    # 展平张量
    flat_tensor = tensor.flatten()
    
    # 计算最小值和最大值
    min_val = torch.min(flat_tensor).item()
    max_val = torch.max(flat_tensor).item()
    
    # 计算箱宽度
    if max_val > min_val:
        bin_width = (max_val - min_val) / bins
        bin_edges = torch.arange(min_val, max_val + bin_width, bin_width)
        
        # 计算直方图
        hist = torch.histc(flat_tensor, bins=bins, min=min_val, max=max_val)
        return hist, bin_edges
    else:
        # 如果最大值等于最小值，返回单一箱
        return torch.tensor([len(flat_tensor)]), torch.tensor([min_val, min_val])

def calculate_statistics(tensor: torch.Tensor) -> Dict[str, float]:
    """
    计算张量的多种统计信息
    
    Args:
        tensor: 输入张量
    
    Returns:
        统计信息字典
    """
    return {
        "norm": calculate_norm(tensor),
        "max": calculate_max(tensor),
        "min": calculate_min(tensor),
        "mean": calculate_mean(tensor),
        "std": calculate_std(tensor),
        "zeros": calculate_zeros(tensor),
        "nans": calculate_nans(tensor),
        "infs": calculate_infs(tensor)
    }

def calculate_sparsity(tensor: torch.Tensor, threshold: float = 1e-6) -> float:
    """
    计算张量的稀疏度（接近零的元素比例）
    
    Args:
        tensor: 输入张量
        threshold: 接近零的阈值
    
    Returns:
        稀疏度（0到1之间）
    """
    total_elements = tensor.numel()
    near_zero_elements = (torch.abs(tensor) < threshold).sum().item()
    return near_zero_elements / total_elements

def calculate_outlier_ratio(tensor: torch.Tensor, std_threshold: float = 3.0) -> float:
    """
    计算张量中异常值的比例（超过均值±std_threshold*标准差的元素）
    
    Args:
        tensor: 输入张量
        std_threshold: 标准差阈值倍数
    
    Returns:
        异常值比例（0到1之间）
    """
    # 展平张量
    flat_tensor = tensor.flatten()
    
    # 计算均值和标准差
    mean = torch.mean(flat_tensor)
    std = torch.std(flat_tensor)
    
    # 计算上下界
    lower_bound = mean - std_threshold * std
    upper_bound = mean + std_threshold * std
    
    # 计算异常值数量
    outliers = ((flat_tensor < lower_bound) | (flat_tensor > upper_bound)).sum().item()
    
    # 计算比例
    total_elements = flat_tensor.numel()
    return outliers / total_elements

def calculate_quantiles(tensor: torch.Tensor, q: List[float] = [0.25, 0.5, 0.75]) -> Dict[str, float]:
    """
    计算张量的分位数
    
    Args:
        tensor: 输入张量
        q: 分位数列表（0到1之间）
    
    Returns:
        分位数字典
    """
    # 展平张量
    flat_tensor = tensor.flatten()
    
    # 计算分位数
    quantiles = {}
    for percentile in q:
        key = f"q{int(percentile * 100)}"
        quantiles[key] = torch.quantile(flat_tensor, percentile).item()
    
    return quantiles 